{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from os.path import join, exists, split\n",
    "import os \n",
    "# os.chdir('/data/xusc/exp/MTPSL')\n",
    "os.chdir(join(os.getcwd(),'..'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = './data/nyuv2_settings/onelabel.pth'\n",
    "import torch \n",
    "labeled = torch.load(path)['labels_weights'].float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "semantic tensor(265.)\n",
      "depth tensor(265.)\n",
      "normal tensor(265.)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "task_num = labeled.shape[1]\n",
    "\n",
    "task_map = {\n",
    "    0: 'semantic',\n",
    "    1: 'depth',\n",
    "    2: 'normal',\n",
    "\n",
    "}\n",
    "for task_id in range(task_num):\n",
    "    print(task_map[task_id], labeled[:,task_id].sum())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import re\n",
    "\n",
    "from jupyters.utils import  * \n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "import torchvision\n",
    "import cv2\n",
    "from dataset.nyuv2ssl import *\n",
    "\n",
    "# nyuv2_train_set = NYUv2_crop(root= join(os.getcwd(),'data/nyuv2'), train=True, augmentation=True, aug_twice=True)\n",
    "nyuv2_train_set = NYUv2_crop(root= join(os.getcwd(),'data/nyuv2'), train=True, augmentation=True, aug_twice=True)\n",
    "nyuv2_test_set = NYUv2(root= join(os.getcwd(),'data/nyuv2'), train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_labels = []\n",
    "for idx, (image, semantic, depth, normal)  in enumerate(nyuv2_test_set):\n",
    "    all_labels.append(semantic.unique().flatten())\n",
    "    \n",
    "    if idx == 100:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.,  0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12.])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cat(all_labels).unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "save_dir = 'logs/debug_aug'\n",
    "make_dir(save_dir)\n",
    "alpha = 0.5\n",
    "\n",
    "for idx in range(nyuv2_train_set.__len__()):\n",
    "\n",
    "\n",
    "    train_data, train_label, train_depth, train_normal, index, image, semantic, \\\n",
    "                depth, normal, trans_params  = nyuv2_train_set[idx]\n",
    "    # image, semantic, depth, normal, index = nyuv2_train_set[idx]\n",
    "\n",
    "    draw_semantics(image, semantic, alpha).save(join(save_dir,\"%06d#semantic.png\"%(idx)))\n",
    "\n",
    "    draw_depth(image, depth, alpha).save(join(save_dir,\"%06d#depth.png\"%(idx)))\n",
    "\n",
    "    draw_normal(image,normal, alpha).save(join(save_dir,\"%06d#normal.png\"%(idx)))\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "192"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transform_semantic_to_colorful_semantics(semantic).max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([288, 384])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "semantic.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mtpsl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
